{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b8608db5",
   "metadata": {},
   "source": [
    "# Run Forecasts of the diffusion models\n",
    "\n",
    "author: Randy Chase <br>\n",
    "email: randy 'dot' chase 'at' colostate.edu \n",
    "\n",
    "This notebook is to run a bunch of forecasts. Please make sure you have the datasets ready for this. If you are using the results of the paper, they are located on dryad!\n",
    "\n",
    "### 1) Load some packages and the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d69f7180",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "#imports \n",
    "import zarr\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from diffusers import UNet2DModel\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "import gc\n",
    "import os \n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "import os\n",
    "import math\n",
    "from typing import List, Optional, Tuple, Union\n",
    "from diffusers.utils.torch_utils import randn_tensor\n",
    "import tqdm \n",
    "import time as timer \n",
    "from diffusers import AutoencoderKL\n",
    "\n",
    "\n",
    "class ZarrDataset(Dataset):\n",
    "    \"\"\"Class to load the zarr data for pytorch. Might need to be optimized if the data are being lazily loaded...\"\"\"\n",
    "    def __init__(self, zarr_store):\n",
    "        self.store = zarr_store\n",
    "        self.data = zarr.open(self.store, mode='r')\n",
    "        self.length = self.data['input_images'].shape[0]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # Load data lazily\n",
    "        input_image = self.data['input_images'][idx]\n",
    "        output_image = self.data['output_images'][idx]\n",
    "        return torch.tensor(output_image, dtype=torch.float16),torch.tensor(input_image, dtype=torch.float16)\n",
    "    \n",
    "def to_K(images):\n",
    "    \"\"\"function to convert from the mean 0, std 1 space back to brightness temp\"\"\"\n",
    "    mean_data=279.0699458792467\n",
    "    std_data =19.32967519050003\n",
    "    return (images * std_data) + mean_data\n",
    "\n",
    "# Initialize the main dataset, choose the: \n",
    "#training dataset\n",
    "# zarr_store = '/mnt/data1/rchas1/TRANSITION/datasets/edm_GOES_ch13_training_dataset.zarr'\n",
    "#validation dataset\n",
    "# zarr_store = '/mnt/data1/rchas1/TRANSITION/datasets/edm_GOES_ch13_validation_dataset.zarr'\n",
    "#test dataset\n",
    "zarr_store = './datasets/edm_GOES_ch13_test_dataset.zarr'\n",
    "dataset = ZarrDataset(zarr_store)\n",
    "\n",
    "#this is the max size that works with generating 10 ensemble members on the GH200, might need to change depneding on the GPU. \n",
    "batch_size = 10\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "#load the latent dataset (data were compressed with the radames VAE to same time)\n",
    "#validation\n",
    "# zarr_store = '/mnt/data1/rchas1/diffusion_10_4_2inputs_2024validation_latent_radames_v3.zarr/'\n",
    "#test\n",
    "zarr_store = './datasets/edm_GOES_ch13_test_dataset_latent.zarr'\n",
    "\n",
    "dataset = ZarrDataset(zarr_store)\n",
    "latent_train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e2aa850",
   "metadata": {},
   "source": [
    "### 2) Get the classes \n",
    "These are the usual classes originally from NVIDIA, but adapted to take our condition and to go with the paper. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eeca7f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EDMPrecond(torch.nn.Module):\n",
    "    \"\"\" Original Func:: https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L519\n",
    "    \n",
    "    This is a wrapper for your diffusers model. It's purpose is to apply the preconditioning that is talked about in Karras et al. (2022)'s EDM paper. \n",
    "    \n",
    "    I've made some changes for the sake of conditional-EDM (the original paper is unconditional).\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "        generation_channels,                # number of channels you want to generate\n",
    "        model,                              # pytorch model from diffusers \n",
    "        use_fp16        = False,             # Execute the underlying model at FP16 precision?\n",
    "        sigma_min       = 0,                # Minimum supported noise level.\n",
    "        sigma_max       = float('inf'),     # Maximum supported noise level.\n",
    "        sigma_data      = 0.5,              # Expected standard deviation of the training data. this was the default from above\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.generation_channels = generation_channels\n",
    "        self.model = model\n",
    "        self.use_fp16 = use_fp16\n",
    "        self.sigma_min = sigma_min\n",
    "        self.sigma_max = sigma_max\n",
    "        self.sigma_data = sigma_data\n",
    "        \n",
    "    def forward(self, x, sigma, force_fp32=False, **model_kwargs):\n",
    "        \n",
    "        \"\"\" \n",
    "        \n",
    "        This method is to 'call' the neural net. But this is the preconditioning from the Karras EDM paper. \n",
    "        \n",
    "        note for conditional, it expects x to have the condition in the channel dim (dim=1). and the images you want to generate should already have noise.\n",
    "        \n",
    "        x: input stacked image with the generation images stacked with the condition images [batch,generation_channels + condition_channels,nx,ny]\n",
    "        sigma: the noise level of the images in batch [??]\n",
    "        force_fp32: this is forcing calculations to be a certain percision. \n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        #for the calculations, use float 32\n",
    "        x = x.to(torch.float32)\n",
    "        #reshape sigma from _ to _ \n",
    "        sigma = sigma.to(torch.float32).reshape(-1, 1, 1, 1)\n",
    "        \n",
    "        #forcing dtype matching\n",
    "        dtype = torch.float16 if (self.use_fp16 and not force_fp32 and x.device.type == 'cuda') else torch.float32\n",
    "        \n",
    "        #get weights from EDM \n",
    "        c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)\n",
    "        c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt()\n",
    "        c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt()\n",
    "        c_noise = sigma.log() / 4\n",
    "\n",
    "        # split out the images you want to generate and the condition, because the scaling will depend on this. \n",
    "        x_noisy = torch.clone(x[:,0:self.generation_channels])\n",
    "        \n",
    "        #the condition\n",
    "        x_condition = torch.clone(x[:,self.generation_channels:])\n",
    "\n",
    "        \n",
    "        #concatinate back with the scaling applied to only the the generation dimension (x_noisy)\n",
    "        model_input_images = torch.cat([x_noisy*c_in, x_condition], dim=1)\n",
    "        \n",
    "        #denoise the image (e.g., run it through your diffusers model) \n",
    "        F_x = self.model((model_input_images).to(dtype), c_noise.flatten(), return_dict=False)[0]\n",
    "        \n",
    "        #force dtype\n",
    "        assert F_x.dtype == dtype\n",
    "        \n",
    "        #apply additional scalings: make sure you apply skip just to the generation dim (x[:,0:generation_channel]) and NOT applied to (x*c_in)\n",
    "        D_x = c_skip * x_noisy + c_out * F_x.to(torch.float32)\n",
    "        \n",
    "        return D_x\n",
    "\n",
    "    def round_sigma(self, sigma):\n",
    "        return torch.as_tensor(sigma)\n",
    "    \n",
    "class EDMPrecond_TF(torch.nn.Module):\n",
    "    \"\"\" Original Func:: https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/networks.py#L519\n",
    "    \n",
    "    This is a wrapper for your diffusers model. It's purpose is to apply the preconditioning that is talked about in Karras et al. (2022)'s EDM paper. \n",
    "    \n",
    "    I've made some changes for the sake of conditional-EDM (the original paper is unconditional).\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "        generation_channels,                # number of channels you want to generate\n",
    "        model,                              # pytorch model from diffusers \n",
    "        use_fp16        = False,             # Execute the underlying model at FP16 precision?\n",
    "        sigma_min       = 0,                # Minimum supported noise level.\n",
    "        sigma_max       = float('inf'),     # Maximum supported noise level.\n",
    "        sigma_data      = 0.5,              # Expected standard deviation of the training data. this was the default from above\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.generation_channels = generation_channels\n",
    "        self.model = model\n",
    "        self.use_fp16 = use_fp16\n",
    "        self.sigma_min = sigma_min\n",
    "        self.sigma_max = sigma_max\n",
    "        self.sigma_data = sigma_data\n",
    "        \n",
    "    def forward(self, x, sigma, force_fp32=False, **model_kwargs):\n",
    "        \n",
    "        \"\"\" \n",
    "        \n",
    "        This method is to 'call' the neural net. But this is the preconditioning from the Karras EDM paper. \n",
    "        \n",
    "        note for conditional, it expects x to have the condition in the channel dim (dim=1). and the images you want to generate should already have noise.\n",
    "        \n",
    "        x: input stacked image with the generation images stacked with the condition images [batch,generation_channels + condition_channels,nx,ny]\n",
    "        sigma: the noise level of the images in batch [??]\n",
    "        force_fp32: this is forcing calculations to be a certain percision. \n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        #for the calculations, use float 32\n",
    "        x = x.to(torch.float32)\n",
    "        #reshape sigma from _ to _ \n",
    "        sigma = sigma.to(torch.float32).reshape(-1, 1, 1, 1)\n",
    "        \n",
    "        #forcing dtype matching\n",
    "        dtype = torch.float16 if (self.use_fp16 and not force_fp32 and x.device.type == 'cuda') else torch.float32\n",
    "        \n",
    "        #get weights from EDM \n",
    "        c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)\n",
    "        c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2).sqrt()\n",
    "        c_in = 1 / (self.sigma_data ** 2 + sigma ** 2).sqrt()\n",
    "        c_noise = sigma.log() / 4\n",
    "\n",
    "        # split out the images you want to generate and the condition, because the scaling will depend on this. \n",
    "        x_noisy = torch.clone(x[:,0:self.generation_channels])\n",
    "        \n",
    "        #the condition\n",
    "        x_condition = torch.clone(x[:,self.generation_channels:])\n",
    "\n",
    "        \n",
    "        #concatinate back with the scaling applied to only the the generation dimension (x_noisy)\n",
    "        model_input_images = torch.cat([x_noisy*c_in, x_condition], dim=1)\n",
    "        \n",
    "        #denoise the image (e.g., run it through your diffusers model) \n",
    "        F_x = self.model((model_input_images).to(dtype), c_noise.flatten(),\n",
    "                         class_labels=torch.zeros(model_input_images.shape[0]).to(torch.int).to(model_input_images.device),\n",
    "                         return_dict=False)[0]\n",
    "        \n",
    "        #force dtype\n",
    "        assert F_x.dtype == dtype\n",
    "        \n",
    "        #apply additional scalings: make sure you apply skip just to the generation dim (x[:,0:generation_channel]) and NOT applied to (x*c_in)\n",
    "        D_x = c_skip * x_noisy + c_out * F_x.to(torch.float32)\n",
    "        \n",
    "        return D_x\n",
    "\n",
    "    def round_sigma(self, sigma):\n",
    "        return torch.as_tensor(sigma)\n",
    "\n",
    "class EDMLoss:\n",
    "    \n",
    "    \"\"\"Original Func:: https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/training/loss.py\n",
    "    \n",
    "    This is the loss function class from Karras et al. (2022)'s EDM paper. Only thing changed here is that the __call__ takes the clean_images and the condition_images seperately. It expects your model to be wrapped with that EDMPrecond class. \n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self, P_mean=-1.2, P_std=1.2, sigma_data=0.5):\n",
    "        \"\"\" These describe the distribution of sigmas we should sample during training \"\"\"\n",
    "        self.P_mean = P_mean\n",
    "        self.P_std = P_std\n",
    "        self.sigma_data = sigma_data\n",
    "\n",
    "    def __call__(self, net, clean_images, condition_images, labels=None, augment_pipe=None):\n",
    "        \n",
    "        \"\"\" \n",
    "        \n",
    "        net: is a pytorch model wrapped with EDMPrecond\n",
    "        clean_images: the images you want to generate, [batch,generation_channels,nx,ny]\n",
    "        condition_images:images you want to condition with [batch,condition_channels,nx,ny]\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        #get random seeds, one for each image in the batch \n",
    "        rnd_normal = torch.randn([clean_images.shape[0], 1, 1, 1], device=clean_images.device)\n",
    "        \n",
    "        #get random noise levels (sigmas)\n",
    "        sigma = (rnd_normal * self.P_std + self.P_mean).exp()\n",
    "        \n",
    "        #get the loss weight for those sigmas \n",
    "        weight = (sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2\n",
    "        \n",
    "        #make the noise scalars images so we can add them to our images\n",
    "        n = torch.randn_like(clean_images) * sigma\n",
    "    \n",
    "        #add noise to the clean images \n",
    "        noisy_images = torch.clone(clean_images + n)\n",
    "        \n",
    "        #cat the images for the wrapped model call \n",
    "        model_input_images = torch.cat([noisy_images, condition_images], dim=1)\n",
    "        \n",
    "        #call the EDMPrecond model \n",
    "        denoised_images = net(model_input_images, sigma)\n",
    "        \n",
    "        #calc the weighted loss at each pixel, the mean across all GPUs and pixels is in the main train_loop \n",
    "        loss = weight * ((denoised_images - clean_images) ** 2)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "class StackedRandomGenerator:  # pragma: no cover\n",
    "    \"\"\"\n",
    "    Wrapper for torch.Generator that allows specifying a different random seed\n",
    "    for each sample in a minibatch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, device, seeds):\n",
    "        super().__init__()\n",
    "        self.generators = [\n",
    "            torch.Generator(device).manual_seed(int(seed) % (1 << 32)) for seed in seeds\n",
    "        ]\n",
    "\n",
    "    def randn(self, size, **kwargs):\n",
    "        if size[0] != len(self.generators):\n",
    "            raise ValueError(\n",
    "                f\"Expected first dimension of size {len(self.generators)}, got {size[0]}\"\n",
    "            )\n",
    "        return torch.stack(\n",
    "            [torch.randn(size[1:], generator=gen, **kwargs) for gen in self.generators]\n",
    "        )\n",
    "\n",
    "    def randn_like(self, input):\n",
    "        return self.randn(\n",
    "            input.shape, dtype=input.dtype, layout=input.layout, device=input.device\n",
    "        )\n",
    "\n",
    "    def randint(self, *args, size, **kwargs):\n",
    "        if size[0] != len(self.generators):\n",
    "            raise ValueError(\n",
    "                f\"Expected first dimension of size {len(self.generators)}, got {size[0]}\"\n",
    "            )\n",
    "        return torch.stack(\n",
    "            [\n",
    "                torch.randint(*args, size=size[1:], generator=gen, **kwargs)\n",
    "                for gen in self.generators\n",
    "            ]\n",
    "        )\n",
    "    \n",
    "def edm_sampler(\n",
    "    net, latents, condition_images,class_labels=None, randn_like=torch.randn_like,\n",
    "    num_steps=18, sigma_min=0.002, sigma_max=80, rho=7,\n",
    "    S_churn=0, S_min=0, S_max=float('inf'), S_noise=1,\n",
    "):\n",
    "    \"\"\" adapted from: https://github.com/NVlabs/edm/blob/008a4e5316c8e3bfe61a62f874bddba254295afb/generate.py \n",
    "    \n",
    "    only thing i had to change was provide a condition as input to this func, then take that input and concat with generated image for the model call. \n",
    "    \n",
    "    \"\"\"\n",
    "    # Adjust noise levels based on what's supported by the network.\n",
    "    sigma_min = max(sigma_min, net.sigma_min)\n",
    "    sigma_max = min(sigma_max, net.sigma_max)\n",
    "\n",
    "    # Time step discretization.\n",
    "    step_indices = torch.arange(num_steps, dtype=torch.float64, device=latents.device)\n",
    "    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho\n",
    "    t_steps = torch.cat([net.round_sigma(t_steps), torch.zeros_like(t_steps[:1])]) # t_N = 0\n",
    "    \n",
    "    # Main sampling loop.\n",
    "    x_next = latents.to(torch.float64) * t_steps[0]\n",
    "    for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1\n",
    "        x_cur = x_next\n",
    "\n",
    "        # Increase noise temporarily.\n",
    "        gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0\n",
    "        t_hat = net.round_sigma(t_cur + gamma * t_cur)\n",
    "        x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * randn_like(x_cur)\n",
    "\n",
    "        #need to concat the condition here \n",
    "        model_input_images = torch.cat([x_hat, condition_images], dim=1)\n",
    "        # Euler step.\n",
    "        with torch.no_grad():\n",
    "            denoised = net(model_input_images, t_hat).to(torch.float64)\n",
    "\n",
    "        d_cur = (x_hat - denoised) / t_hat\n",
    "        x_next = x_hat + (t_next - t_hat) * d_cur\n",
    "\n",
    "        # Apply 2nd order correction.\n",
    "        if i < num_steps - 1:\n",
    "            model_input_images = torch.cat([x_next, condition_images], dim=1)\n",
    "            with torch.no_grad():\n",
    "                denoised = net(model_input_images, t_next).to(torch.float64)\n",
    "            d_prime = (x_next - denoised) / t_next\n",
    "            x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)\n",
    "\n",
    "    return x_next\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88d56c72",
   "metadata": {},
   "source": [
    "### 3) Define wrapper functions\n",
    "\n",
    "These functions help do the forecasts by wrapping them and doing the autoregression for you. Alot is similar between them, but there are some differences with shapes and what not. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c255003f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_forecast_diff_ens(model,images_condition,ens_size=10,same_seed=True,time_steps_forward = np.arange(0,18),S_churn=np.sqrt(2) - 1,num_steps=18,sigma_max=80, rho=7):\n",
    "    \"\"\"\n",
    "    \n",
    "    This is to run an ensemble forecast for the plain EDM diffusion models.\n",
    "    \n",
    "    INPUTS\n",
    "    model: the EDMPrecondition wrapped pytorch model \n",
    "    images_condition: a tensor of shape [batch,condition,nx,ny] containing the condition images for the start of the forecast \n",
    "    ens_size: an integer for the number of ensemble members that are requested. default is 10\n",
    "    same_seed: set the seeds the same for all forecast steps. This is off by default because it breaks the autoregression.\n",
    "    time_steps_forward: array of the forecast length \n",
    "    \n",
    "    # Karras et al. params passed through to the edm_sampler \n",
    "    time_steps_forward: array of the number of steps you want to take forward. by default its 3 hours for our paper. \n",
    "    S_churn: how big of a step backwards you want to take in the SDE image generation. If its 0, this will just turn off. By default I set it to the max value. \n",
    "    num_steps: number of diffusion steps you want to take. \n",
    "    sigma_max: the starting sigma value\n",
    "    rho: how the steps are spaced \n",
    "     \n",
    "    OUTPUTS\n",
    "    full_image_list: full list of forecast, first 2 images were the original inputs. The rest are all the ML forecasts. \n",
    "    \"\"\"\n",
    "\n",
    "    #define arrays to store things as we generate the sequence \n",
    "    full_image_list = torch.zeros([images_condition.shape[0],ens_size,time_steps_forward.shape[0]+2,images_condition.shape[2],images_condition.shape[3]])\n",
    "    \n",
    "    #fill the first two in the full sequence with the condition, make sure you repeat it to match the requested ens size\n",
    "    full_image_list[:,:,0:2] = torch.clone(images_condition.unsqueeze(1)).repeat((1,ens_size,1,1,1)).to('cpu')\n",
    "    \n",
    "    #loop over requested time steps \n",
    "    for time in tqdm.tqdm(time_steps_forward):\n",
    "        \n",
    "        #if you by chance want to use the same noisy seed for every time step, use same_seed. This broke my autoregression before....\n",
    "        if same_seed:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist())\n",
    "        else:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist() + time)\n",
    "        \n",
    "        #build out the random noise images \n",
    "        latents = rnd.randn([full_image_list.shape[0]*ens_size, 1, images_condition.shape[2], images_condition.shape[3]],device='cuda')\n",
    "        \n",
    "        #create the current condition tensor \n",
    "        current_image_list = torch.clone(full_image_list[:,:,time:time+2]).to(torch.float32).to('cuda')\n",
    "\n",
    "        #reshape to run through NN, needs to be [batch*ens_size,condition,nx,ny]\n",
    "        current_image_list = current_image_list.reshape([full_image_list.shape[0]*full_image_list.shape[1],\n",
    "                                                        2,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "\n",
    "        #run diffusion model 10 min forecast \n",
    "        images_ens = edm_sampler(model,latents,current_image_list,num_steps=num_steps,S_churn=S_churn,sigma_max=sigma_max,rho=rho)\n",
    "        \n",
    "        #undo reshape back to [batch,ens_size,1,nx,ny]\n",
    "        images_ens = images_ens.reshape([full_image_list.shape[0],full_image_list.shape[1],1,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "        \n",
    "        #save that 10 min forecast out so you can keep the whole sequence \n",
    "        full_image_list[:,:,time+2:time+3] = torch.clone(images_ens.cpu())\n",
    "             \n",
    "    return full_image_list\n",
    "\n",
    "def run_forecast_corrdiff_ens(model_diff,model_unet,images_condition,same_seed=True,ens_size=10,time_steps_forward = np.arange(0,18),S_churn=np.sqrt(2) - 1,num_steps=18,sigma_max=80, rho=7):\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    This is to run an ensemble forecast for the CorrDiff EDM diffusion models.\n",
    "    \n",
    "    INPUTS\n",
    "    model_diff: the EDMPrecondition wrapped pytorch model \n",
    "    model_unet: the Unet model that was used to get the 'first guess' for the corrdiff approach. \n",
    "    images_condition: a tensor of shape [batch,condition,nx,ny] containing the condition images for the start of the forecast \n",
    "    ens_size: an integer for the number of ensemble members that are requested. default is 10\n",
    "    same_seed: set the seeds the same for all forecast steps. This is off by default because it breaks the autoregression.\n",
    "    time_steps_forward: array of the forecast length \n",
    "    \n",
    "    # Karras et al. params passed through to the edm_sampler \n",
    "    time_steps_forward: array of the number of steps you want to take forward. by default its 3 hours for our paper. \n",
    "    S_churn: how big of a step backwards you want to take in the SDE image generation. If its 0, this will just turn off. By default I set it to the max value. \n",
    "    num_steps: number of diffusion steps you want to take. \n",
    "    sigma_max: the starting sigma value\n",
    "    rho: how the steps are spaced \n",
    "     \n",
    "    OUTPUTS\n",
    "    full_image_list: full list of forecast, first 2 images were the original inputs. The rest are all the ML forecasts. \n",
    "    \"\"\"\n",
    "    #define arrays to store things as we generate the sequence \n",
    "    full_image_list = torch.zeros([images_condition.shape[0],ens_size,time_steps_forward.shape[0]+2,images_condition.shape[2],images_condition.shape[3]])\n",
    "    #fill the first two in the full sequence with the condition, make sure you repeat it to match the requested ens size\n",
    "    full_image_list[:,:,0:2] = torch.clone(images_condition.unsqueeze(1)).repeat((1,ens_size,1,1,1)).to('cpu')\n",
    "    \n",
    "    #loop over requested time steps \n",
    "    for time in tqdm.tqdm(time_steps_forward):\n",
    "        \n",
    "        #if you by chance want to use the same noisy seed for every time step, use same_seed. This broke my autoregression before....\n",
    "        if same_seed:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist())\n",
    "        else:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist() + time)\n",
    "        \n",
    "        #build out the random noise images \n",
    "        latents = rnd.randn([full_image_list.shape[0]*ens_size, 1, images_condition.shape[2], images_condition.shape[3]],device='cuda')\n",
    "        \n",
    "        #create the current condition tensor \n",
    "        current_image_list = torch.clone(full_image_list[:,:,time:time+2]).to(torch.float32).to('cuda')\n",
    "        \n",
    "        #reshape to run through NN, needs to be [batch*ens_size,condition,nx,ny]\n",
    "        current_image_list = current_image_list.reshape([full_image_list.shape[0]*full_image_list.shape[1],\n",
    "                                                        2,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "        \n",
    "        #STEP 1: run unet to get 'blurry' forecast \n",
    "        with torch.no_grad(): #turn off gradient tracking to ensure memory savings, the zeros here are a hack of how i used the same unet. \n",
    "            images_unet = model_unet(current_image_list,torch.zeros(current_image_list.shape[0]).to(current_image_list.device), return_dict=False)[0]\n",
    "        \n",
    "        #concat the unet output to the other current images end shape will be [batch*ens_size,3,nx,ny]\n",
    "        current_image_list_diff = torch.concat([current_image_list,images_unet],axis=1)\n",
    "        \n",
    "        #STEP 2: run diffusion to get residual for adjusting the unet 'blurry' forecast  \n",
    "        images_diff = edm_sampler(model_diff,latents,current_image_list_diff,num_steps=num_steps,S_churn=S_churn,sigma_max=sigma_max,rho=rho)\n",
    "        \n",
    "        #undo the 0-1 scaling I did to help the ML learn \n",
    "        data_mean = torch.tensor(-0.0009)\n",
    "        data_std = torch.tensor(0.0807)\n",
    "        images_diff = images_diff*data_std + data_mean\n",
    "        \n",
    "        #undo reshape back to [batch,ens_size,1,nx,ny]\n",
    "        images_diff =  images_diff.reshape([full_image_list.shape[0],full_image_list.shape[1],1,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "        images_unet =  images_unet.reshape([full_image_list.shape[0],full_image_list.shape[1],1,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "        \n",
    "        #add resid to unet output \n",
    "        images_ens = images_unet.cpu() + images_diff.cpu()\n",
    "        \n",
    "        #store new 'starting' state\n",
    "        full_image_list[:,:,time+2:time+3] = torch.clone(images_ens.cpu())\n",
    "\n",
    "    return full_image_list\n",
    "\n",
    "def run_forecast_ldm_ens(model,images_condition,images_condition_latent,vae,same_seed=True,ens_size=10,time_steps_forward = np.arange(0,18,dtype=int),S_churn=np.sqrt(2) - 1,num_steps=18,sigma_max=80, rho=7):\n",
    "    \n",
    "    \"\"\"\n",
    "    This is to run an ensemble forecast for the latent EDM diffusion models.\n",
    "    \n",
    "    INPUTS\n",
    "    model: the EDMPrecondition wrapped pytorch model \n",
    "    images_condition: a tensor of shape [batch,condition,nx,ny] containing the condition images for the start of the forecast\n",
    "    images_condition_latent: the same tensor as before, but pre-compressed through the VAE to save time. \n",
    "    vae: the huggingface pretrained pytorch VAE.\n",
    "    ens_size: an integer for the number of ensemble members that are requested. default is 10\n",
    "    same_seed: set the seeds the same for all forecast steps. This is off by default because it breaks the autoregression.\n",
    "    time_steps_forward: array of the forecast length \n",
    "    \n",
    "    # Karras et al. params passed through to the edm_sampler \n",
    "    time_steps_forward: array of the number of steps you want to take forward. by default its 3 hours for our paper. \n",
    "    S_churn: how big of a step backwards you want to take in the SDE image generation. If its 0, this will just turn off. By default I set it to the max value. \n",
    "    num_steps: number of diffusion steps you want to take. \n",
    "    sigma_max: the starting sigma value\n",
    "    rho: how the steps are spaced \n",
    "     \n",
    "    OUTPUTS\n",
    "    full_image_list: full list of forecast, first 2 images were the original inputs. The rest are all the ML forecasts. \n",
    "    \"\"\"\n",
    "    \n",
    "    #define arrays to store things as we generate the sequence \n",
    "    full_image_list = torch.zeros([images_condition.shape[0],ens_size,time_steps_forward.shape[0]+2,images_condition.shape[2],images_condition.shape[3]])\n",
    "    full_image_list[:,:,0:2] = torch.clone(images_condition.unsqueeze(1)).repeat((1,ens_size,1,1,1)).to('cpu')\n",
    "    \n",
    "    #define arrays to store things as we generate the sequence but also for the latent space to keep track there, \n",
    "    #the 8 and the 4 come from the added 4 channels in the latent space for each image in the full space \n",
    "    latent_image_list = torch.zeros([images_condition.shape[0],ens_size, 8 + time_steps_forward.shape[0]*4,images_condition_latent.shape[2],images_condition_latent.shape[3]])\n",
    "    latent_image_list[:,:,0:8] = torch.clone(images_condition_latent.unsqueeze(1)).repeat((1,ens_size,1,1,1)).to('cpu')\n",
    "    \n",
    "    #these are scalars from the training set to undo the scaling of the latent space that will then be fed into the vae.\n",
    "    condition_std = torch.tensor(6.0271)\n",
    "    condition_mean = torch.tensor(-4.2534)\n",
    "\n",
    "    for time in tqdm.tqdm(time_steps_forward):\n",
    "        \n",
    "        #if you by chance want to use the same noisy seed for every time step, use same_seed. This broke my autoregression before....\n",
    "        if same_seed:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist())\n",
    "        else:\n",
    "            rnd = StackedRandomGenerator('cuda',np.arange(0,full_image_list.shape[0]*ens_size,1).astype(int).tolist() + time)\n",
    "         \n",
    "        #build out the random noise images \n",
    "        latents = rnd.randn([images_condition.shape[0]*ens_size, 4, latent_image_list.shape[3], latent_image_list.shape[4]],device='cuda')\n",
    "        \n",
    "        #create the current condition tensor (again the 4 and 8 come from the 4 extra channels in the latent space)\n",
    "        current_image_list = torch.clone(latent_image_list[:,:,(time*4):(time*4)+8]).to('cuda')\n",
    "        \n",
    "        #reshape to run through NN [batch*ens_size,8,nx,ny]\n",
    "        current_image_list = current_image_list.reshape([latent_image_list.shape[0]*latent_image_list.shape[1],\n",
    "                                                        8,latent_image_list.shape[3],latent_image_list.shape[4]])\n",
    "\n",
    "        #run diffusion model 10 min forecast \n",
    "        images_batch = edm_sampler(model,latents,current_image_list,num_steps=num_steps,S_churn=S_churn,sigma_max=sigma_max,rho=rho)\n",
    "        \n",
    "        #store latents for the next time step, reshape it to be [batch,ens_size,nx,ny]\n",
    "        latent_image_list[:,:,(time*4+8):(time*4+12)]  = torch.clone(images_batch.reshape([latent_image_list.shape[0],latent_image_list.shape[1],4,latent_image_list.shape[3],latent_image_list.shape[4]]).cpu())\n",
    "\n",
    "        #undo mean0,std1 scalaing before decoding (done to make training the denoiser easier)\n",
    "        images_batch = (images_batch*condition_std) + condition_mean\n",
    "        \n",
    "        #decode with the pretrained vae \n",
    "        with torch.no_grad():\n",
    "            reconstructed_images_batch = vae.decode(images_batch.to(torch.float)) #input should be [batch*channel,4,64,64]\n",
    "            #take the mean across RGB, this is needed because we are using a vae for RGB images \n",
    "            reconstructed_images_batch = reconstructed_images_batch.sample.mean(axis=1).unsqueeze(1) #reconstructed should be [batch*ens_size,3,nx,ny]\n",
    "    \n",
    "        #undo 0-1 scaling that was used for the VAE \n",
    "        data_min = -6.16015625\n",
    "        data_max = 3.15234375\n",
    "        reconstructed_images_batch = reconstructed_images_batch*(data_max-data_min) + data_min\n",
    "        \n",
    "        #reshape to [batch,ens_size,nx,ny]\n",
    "        reconstructed_images_batch = reconstructed_images_batch.reshape([full_image_list.shape[0],full_image_list.shape[1],1,full_image_list.shape[3],full_image_list.shape[4]])\n",
    "        \n",
    "        #store all the forecast times \n",
    "        full_image_list[:,:,time+2:time+3] = torch.clone(reconstructed_images_batch.cpu())\n",
    "        \n",
    "    return full_image_list\n",
    "\n",
    "def run_forecast_unet(model,images_condition,time_steps_forward = np.arange(0,18)):\n",
    "    \"\"\"\n",
    "    This is to run a vanilla unet forecast \n",
    "    \n",
    "    INPUTS\n",
    "    model: the pytorch unet, from diffusers here \n",
    "    images_condition: a tensor of shape [batch,condition,nx,ny] containing the condition images for the start of the forecast\n",
    "    time_steps_forward: array of the forecast length \n",
    "     \n",
    "    OUTPUTS\n",
    "    full_image_list: full list of forecast, first 2 images were the original inputs. The rest are all the ML forecasts. \n",
    "    \"\"\"\n",
    "\n",
    "    full_image_list = torch.zeros([images_condition.shape[0],time_steps_forward.shape[0]+2,images_condition.shape[2],images_condition.shape[3]])\n",
    "    full_image_list[:,0:2] = torch.clone(images_condition).to('cpu')\n",
    "\n",
    "    for time in tqdm.tqdm(time_steps_forward):\n",
    "        current_image_list = torch.clone(full_image_list[:,time:time+2]).to(torch.float32).to('cuda')\n",
    "\n",
    "        with torch.no_grad():\n",
    "            images_ens = model(current_image_list,torch.zeros(current_image_list.shape[0]).to(current_image_list.device), return_dict=False)[0]\n",
    "\n",
    "        full_image_list[:,time+2:time+3] = torch.clone(images_ens.cpu())\n",
    "        \n",
    "    return full_image_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4af9806",
   "metadata": {},
   "source": [
    "### 4) Load in the models!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d9004c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "###################################\n",
    "###### Load plain EDM model #######\n",
    "###################################\n",
    "\n",
    "#build arch \n",
    "model = UNet2DModel(\n",
    "    sample_size=256,  # the target image resolution\n",
    "    in_channels=3,  # the number of input channels, noisy + 2 condition channels \n",
    "    out_channels=1,  # the number of output channels, just 1 10 min forecast \n",
    "    layers_per_block=2,  # how many ResNet layers to use per UNet block\n",
    "    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n",
    "    down_block_types=(\n",
    "        \"DownBlock2D\",  # a regular ResNet downsampling block\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n",
    "        \"DownBlock2D\",\n",
    "    ),\n",
    "    up_block_types=(\n",
    "        \"UpBlock2D\",  # a regular ResNet upsampling block\n",
    "        \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "    ),\n",
    ")\n",
    "#get weights \n",
    "checkpoint = torch.load(\"./models/edm_plain_diffusion/checkpoint.pth\", map_location='cuda')\n",
    "#load weights \n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "#put it on the gpu \n",
    "model.to('cuda')\n",
    "#wrap it with the precondition\n",
    "model_wrapped = EDMPrecond(1,model)\n",
    "\n",
    "########################################\n",
    "######### Load CorrDiff model ##########\n",
    "########################################\n",
    "\n",
    "#build arch \n",
    "model_corrdiff = UNet2DModel(\n",
    "    sample_size=256,  # the target image resolution\n",
    "    in_channels=4,  # the number of input channels, noisy + 2 conidtion + 1 unet forecast\n",
    "    out_channels=1,  # the number of output channels\n",
    "    layers_per_block=2,  # how many ResNet layers to use per UNet block\n",
    "    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n",
    "    down_block_types=(\n",
    "        \"DownBlock2D\",  # a regular ResNet downsampling block\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n",
    "        \"DownBlock2D\",\n",
    "    ),\n",
    "    up_block_types=(\n",
    "        \"UpBlock2D\",  # a regular ResNet upsampling block\n",
    "        \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "    ),\n",
    ")\n",
    "#get weights\n",
    "checkpoint = torch.load(\"./models/edm_corrdiff/checkpoint.pth\", map_location='cuda')\n",
    "#load weights \n",
    "model_corrdiff.load_state_dict(checkpoint['model_state_dict'])\n",
    "#put it on the gpu\n",
    "model_corrdiff.to('cuda')\n",
    "#wrap it with the precondition\n",
    "model_corrdiff_wrapped = EDMPrecond(1,model_corrdiff)\n",
    "\n",
    "\n",
    "########################################\n",
    "############ Load LDM model ############\n",
    "########################################\n",
    "\n",
    "#build arch \n",
    "model_ldm = UNet2DModel(\n",
    "    sample_size=64,  # the target image resolution\n",
    "    in_channels=12,  # Expands alot because the latent space adds 4 channels per one input, (4*3)\n",
    "    out_channels=4,  # the number of output channels\n",
    "    layers_per_block=2,  # how many ResNet layers to use per UNet block\n",
    "    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n",
    "    down_block_types=(\n",
    "        \"DownBlock2D\",  # a regular ResNet downsampling block\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"DownBlock2D\",\n",
    "        \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n",
    "        \"DownBlock2D\",\n",
    "    ),\n",
    "    up_block_types=(\n",
    "        \"UpBlock2D\",  # a regular ResNet upsampling block\n",
    "        \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "        \"UpBlock2D\",\n",
    "    ),\n",
    ")\n",
    "#get weights\n",
    "checkpoint = torch.load(\"./models/edm_ldm_radames/checkpoint.pth\", map_location='cuda')\n",
    "#load weights \n",
    "model_ldm.load_state_dict(checkpoint['model_state_dict'])\n",
    "#put it on the gpu\n",
    "model_ldm.to('cuda')\n",
    "#wrap it with the precondition\n",
    "model_ldm_wrapped = EDMPrecond(4,model_ldm)\n",
    "\n",
    "########################################\n",
    "######## Load plain unet model #########\n",
    "########################################\n",
    "\n",
    "model_unet = UNet2DModel.from_pretrained(\"./models/unet_vanilla/\").to('cuda')\n",
    "\n",
    "########################################\n",
    "######### Load pretrained vae ##########\n",
    "########################################\n",
    "\n",
    "#this is the path on hugging face, not locally. \n",
    "vae =  AutoencoderKL.from_pretrained(\"radames/stable-diffusion-x4-upscaler-img2img\", subfolder=\"vae\").to('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1fca09e",
   "metadata": {},
   "source": [
    "### 5) Run the forecast \n",
    "\n",
    "This loop will run the forecast for all methods, the plain EDM first, then the CorrDiff, then the unet and LDM. It also passes along a persistence forecast and a the truth. \n",
    "\n",
    "This takes the longest amount of time. With the generation params I have selected, its about 10 mins for each of the big ones. \n",
    "\n",
    "11 mins for the plain edm diffusion \n",
    "11 mins for the corrdiff \n",
    " 0 for the vanilla unet \n",
    " 1 for the LDM forecast \n",
    "------------------------\n",
    "23 mins total per batch size of 10\n",
    "\n",
    "1024 total images \n",
    "\n",
    "~103 batches \n",
    "\n",
    "(23*100)/60  ~= 40 hours to run... come back later if this is in a screen "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "076fb853",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]\n",
      "  0%|                                                                                                                                                     | 0/18 [00:00<?, ?it/s]\u001b[A\n",
      "  6%|███████▊                                                                                                                                     | 1/18 [01:19<22:39, 80.00s/it]\u001b[A\n",
      " 11%|███████████████▋                                                                                                                             | 2/18 [02:39<21:17, 79.86s/it]\u001b[A\n",
      " 17%|███████████████████████▌                                                                                                                     | 3/18 [03:59<19:57, 79.83s/it]\u001b[A\n",
      " 22%|███████████████████████████████▎                                                                                                             | 4/18 [05:19<18:37, 79.81s/it]\u001b[A\n",
      " 28%|███████████████████████████████████████▏                                                                                                     | 5/18 [06:39<17:17, 79.80s/it]\u001b[A\n",
      " 33%|███████████████████████████████████████████████                                                                                              | 6/18 [07:58<15:57, 79.80s/it]\u001b[A\n",
      " 39%|██████████████████████████████████████████████████████▊                                                                                      | 7/18 [09:18<14:37, 79.77s/it]\u001b[A\n",
      " 44%|██████████████████████████████████████████████████████████████▋                                                                              | 8/18 [10:38<13:17, 79.72s/it]\u001b[A\n",
      " 50%|██████████████████████████████████████████████████████████████████████▌                                                                      | 9/18 [11:57<11:57, 79.68s/it]\u001b[A"
     ]
    }
   ],
   "source": [
    "#3 hour forecast, this can be changed to your liking. \n",
    "time_steps_forward = np.arange(0,18)\n",
    "time_len = time_steps_forward.shape[0] + 2\n",
    "\n",
    "# some image shapes for all the arrays (for validation and test)\n",
    "nx = 256\n",
    "ens_size = 10\n",
    "num_forecasts = 1000\n",
    "\n",
    "\n",
    "#preallocate arrays so we can store the data because we only gen a batch at a time \n",
    "full_image_list_K_concat_diff_ens_mems = torch.zeros([num_forecasts,ens_size,time_len,nx,nx])\n",
    "full_image_list_K_concat_corrdiff_ens_mems = torch.zeros([num_forecasts,ens_size,time_len,nx,nx])\n",
    "full_image_list_K_concat_ldm_ens_mems = torch.zeros([num_forecasts,ens_size,time_len,nx,nx])\n",
    "full_image_list_K_concat_unet = torch.zeros([num_forecasts,time_len,nx,nx])\n",
    "#these are hardcoded because these are fixed \n",
    "full_image_list_K_concat_pers = torch.zeros([num_forecasts,1,nx,nx])\n",
    "clean_image_list_K_concat = torch.zeros([num_forecasts,18,nx,nx])\n",
    "\n",
    "#best set of Karras 2022 generation parameters for the latent model, assuming it works for all of them. \n",
    "num_steps = 36\n",
    "S_churn=7.2\n",
    "sigma_max = 140 \n",
    "rho = 4\n",
    "\n",
    "#Karras 2022 defaults\n",
    "# num_steps = 18\n",
    "# S_churn=0\n",
    "# sigma_max = 80\n",
    "# rho = 7\n",
    "\n",
    "#main loop, for each batch, run the forecast \n",
    "for step, (batch, batch_latent) in enumerate(tqdm.tqdm(zip(train_dataloader, latent_train_dataloader))):\n",
    "        \n",
    "        #these if's are if you need to start stop from a specific batch \n",
    "        if (step >= 1) and (step < 2):\n",
    "            #the true data \n",
    "            clean_images_eval = batch[0]\n",
    "            #the condition, throw it on the GPU cause we need it \n",
    "            condition_images_eval = batch[1].to('cuda')\n",
    "            \n",
    "            #the true latent image \n",
    "            latent_clean_images_eval = batch_latent[0]\n",
    "            #the latent condition , throw it on the GPU cause we need it \n",
    "            latent_condition_images_eval = batch_latent[1].to('cuda')\n",
    "            \n",
    "            #run forecasts \n",
    "            full_image_list_diff_mems_ens = run_forecast_diff_ens(model_wrapped,condition_images_eval,S_churn=S_churn,\n",
    "                                                                num_steps=num_steps,rho=rho,sigma_max=sigma_max,same_seed=False,time_steps_forward=time_steps_forward)\n",
    "            full_image_list_corrdiff_ens_mems = run_forecast_corrdiff_ens(model_corrdiff_wrapped,model_unet,condition_images_eval,S_churn=S_churn,\n",
    "                                                                num_steps=num_steps,rho=rho,sigma_max=sigma_max,same_seed=False,time_steps_forward=time_steps_forward)\n",
    "            full_image_list_unet = run_forecast_unet(model_unet,condition_images_eval,time_steps_forward=time_steps_forward)\n",
    "            full_image_list_ldm_ens_mems = run_forecast_ldm_ens(model_ldm_wrapped,condition_images_eval,latent_condition_images_eval,vae,S_churn=S_churn,\n",
    "                                                                    num_steps=num_steps,rho=rho,sigma_max=sigma_max,same_seed=False,time_steps_forward=time_steps_forward)\n",
    "\n",
    "            #convert back to IR brightness temperatures \n",
    "            full_image_list_diff_ens_mems_K = to_K(full_image_list_diff_mems_ens)\n",
    "            full_image_list_corrdiff_ens_mems_K = to_K(full_image_list_corrdiff_ens_mems)\n",
    "            full_image_list_ldm_ens_mems_K = to_K(full_image_list_ldm_ens_mems)\n",
    "            full_image_list_unet_K = to_K(full_image_list_unet)\n",
    "            clean_images_eval_K = to_K(clean_images_eval)\n",
    "            persistence_images_K = to_K(condition_images_eval[:,1:2].cpu())\n",
    "            \n",
    "            #store them \n",
    "            full_image_list_K_concat_diff_ens_mems[step*batch_size:(step+1)*batch_size] = full_image_list_diff_ens_mems_K\n",
    "            full_image_list_K_concat_corrdiff_ens_mems[step*batch_size:(step+1)*batch_size] = full_image_list_corrdiff_ens_mems_K\n",
    "            full_image_list_K_concat_unet[step*batch_size:(step+1)*batch_size] = full_image_list_unet_K\n",
    "            clean_image_list_K_concat[step*batch_size:(step+1)*batch_size] = clean_images_eval_K\n",
    "            full_image_list_K_concat_pers[step*batch_size:(step+1)*batch_size] = persistence_images_K\n",
    "            full_image_list_K_concat_ldm_ens_mems[step*batch_size:(step+1)*batch_size] = full_image_list_ldm_ens_mems_K\n",
    "        elif step < 1:\n",
    "            continue \n",
    "        else:\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bab6431d",
   "metadata": {},
   "source": [
    "### 6) Save the forecasts to disk \n",
    "\n",
    "I use xarray, you could use whatever you want, I am only saving out the forecasts, not the inputs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02238540",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "ds_forecasts = xr.Dataset({'corr_diff_mems':(['n_sample','n_member','forecast_time','nx','ny'],full_image_list_K_concat_corrdiff_ens_mems[:,:,2:]),\n",
    "           'diff_mems':(['n_sample','n_member','forecast_time','nx','ny'],full_image_list_K_concat_diff_ens_mems[:,:,2:]),\n",
    "           'truth':(['n_sample','forecast_time','nx','ny'],clean_image_list_K_concat[:,:]),\n",
    "           'unet':(['n_sample','forecast_time','nx','ny'],full_image_list_K_concat_unet[:,2:]),\n",
    "           'ldm_mems':(['n_sample','n_member','forecast_time','nx','ny'],full_image_list_K_concat_ldm_ens_mems[:,:,2:]),\n",
    "            'pers':(['n_sample','nx','ny'],full_image_list_K_concat_pers[:,0])})\n",
    "\n",
    "# Specify compression for each variable\n",
    "encoding = {\n",
    "    'corr_diff_mems': {'zlib': True, 'complevel': 5},\n",
    "    'diff_mems': {'zlib': True, 'complevel': 5},\n",
    "    'ldm_mems': {'zlib': True, 'complevel': 5},\n",
    "    'truth': {'zlib': True, 'complevel': 5},\n",
    "    'unet': {'zlib': True, 'complevel': 5},\n",
    "    'pers': {'zlib': True, 'complevel': 5},\n",
    "}\n",
    "\n",
    "ds_forecasts.to_netcdf('./datasets/test_set_forecasts.nc', encoding=encoding)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
